#include "thinning.h"
#include "thinning_kernel.cuh"
#include "stopwatch.hpp"

extern "C" void initSkeletoniser(Skeletoniser *skel) {
  skel->coordinateListWork_dev=NULL;
  skel->coordinateList_dev=NULL;
  skel->subListLengths_dev=NULL;
  skel->subListLengths=NULL;
  skel->nonZeroSubListLengths=NULL;
  skel->yzPlane_dev=NULL;
  skel->yzPlane=NULL; //for testing purposes only
  skel->voxels_dev=NULL;
  skel->gather_dev=NULL;
  skel->gatherMemSizes=NULL;
  skel->coordinateList_dev=NULL;
  skel->yzPlane_dev=NULL;
}

extern "C" int setupSkeletoniser(Skeletoniser *skel,const dim3 volDim,const int numVoxels) {
  skel->numVoxels=numVoxels;
  skel->volDim=volDim;
  skel->planeDim.x=volDim.y;
  skel->planeDim.y=volDim.z;

  skel->voxels_dev=NULL;
  skel->gather_dev=NULL;
  skel->gatherMemSizes=NULL;
  skel->coordinateList_dev=NULL;
  skel->yzPlane_dev=NULL;


  cudaError_t error;
  cudaDeviceProp prop;
  int numDevices;
  error=cudaGetDeviceCount(&numDevices);
  if (error!=cudaSuccess) {
    printf("CUDA error (could not get device count): %s\n",cudaGetErrorString(error));
    return 1;
  }
  if (numDevices<=0) {
    printf("Could not find any CUDA capable devices.\n");
    return 1;
  }
  error=cudaGetDeviceProperties(&prop,0);
  if (error!=cudaSuccess) {
    printf("CUDA error (could not get device properties): %s\n",cudaGetErrorString(error));
    return 1;
  }
  cudaSetDevice(numDevices-1);
  int maxBlockSize=16;
  if (prop.regsPerBlock/(maxBlockSize*maxBlockSize)<64) {
    maxBlockSize=8;
  }

  /* set up execution configuration */
  int minDim=skel->volDim.x<skel->volDim.y?skel->volDim.x:skel->volDim.y;
  minDim=minDim<skel->volDim.z?minDim:skel->volDim.z;
  skel->blockSize=min(minDim,maxBlockSize);
  skel->numThreads=skel->blockSize*skel->blockSize;
  skel->oneDimGrid.x=skel->numVoxels/skel->numThreads;
  skel->oneDimBlock.x=skel->numThreads;
  skel->twoDimGrid.x=skel->volDim.x/skel->blockSize;
  skel->twoDimGrid.y=skel->volDim.y*skel->volDim.z/skel->blockSize;
  skel->twoDimBlock.x=skel->blockSize;
  skel->twoDimBlock.y=skel->blockSize;
  skel->findCoordinatesGrid.x=skel->planeDim.x*skel->planeDim.y/skel->numThreads;

  /* set up device memory */
  skel->voxelMemSize=sizeof(unsigned char)*skel->numVoxels;

  error=cudaMalloc((void **) &(skel->voxels_dev),skel->voxelMemSize);
  if (error!=cudaSuccess) {
    printf("CUDA error (could not allocate memory for voxel data): %s\n",cudaGetErrorString(error));
    cleanUp(skel);
    return 1;
  }

  /* set up device memory for gather operation */
  skel->gatherLevels=(int)ceil(log(skel->numVoxels)/log(skel->numThreads));
  skel->gather_dev=(unsigned char **)malloc((skel->gatherLevels+1)*sizeof(unsigned char *));
  skel->gatherMemSizes=(unsigned int *)malloc((skel->gatherLevels+1)*sizeof(int));
  int tmp=skel->numVoxels;
  for (int i=0;i<=skel->gatherLevels;i++) {
    unsigned int gatherMemSize=sizeof(unsigned char)*tmp;
    skel->gatherMemSizes[i]=gatherMemSize;
    skel->gather_dev[i]=NULL;
    error=cudaMalloc((void **) &(skel->gather_dev[i]),gatherMemSize);
    if (error!=cudaSuccess) {
      printf("CUDA error (could not allocate memory for gathered data): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
    tmp/=skel->numThreads;
    if (tmp!=1 && tmp<skel->numThreads) tmp=skel->numThreads;
  }
  skel->gathered=(unsigned char *)malloc(skel->gatherMemSizes[1]);

  /* set up device memory for projection (y-dim * z-dim uchar) and strip coordinate list (3 * y-dim * z-dim ushort) */
  skel->coordinateListMemSize=3*volDim.y*volDim.z*sizeof(unsigned short);
  error=cudaMalloc((void **) &(skel->coordinateListWork_dev),skel->coordinateListMemSize);
  if (error!=cudaSuccess) {
    printf("CUDA error (could not allocate memory for strip coordinate list working memory): %s\n",cudaGetErrorString(error));
    cleanUp(skel);
    return 1;
  }

  skel->yzPlaneMemSize=volDim.y*volDim.z*sizeof(unsigned char);
  skel->yzPlane=(unsigned char *)malloc(skel->yzPlaneMemSize);
  error=cudaMalloc((void **) &(skel->yzPlane_dev),skel->yzPlaneMemSize);
  if (error!=cudaSuccess) {
    printf("CUDA error (could not allocate memory for yz plane): %s\n",cudaGetErrorString(error));
    cleanUp(skel);
    return 1;
  }

  skel->subListLengthsMemSize=skel->planeDim.x*skel->planeDim.y/skel->numThreads*sizeof(unsigned short);
  skel->subListLengths=(unsigned short *)malloc(skel->subListLengthsMemSize);
  skel->nonZeroSubListLengths=(unsigned short *)malloc(skel->subListLengthsMemSize);
  error=cudaMalloc((void **) &(skel->subListLengths_dev),skel->subListLengthsMemSize);
  if (error!=cudaSuccess) {
    printf("CUDA error (could not allocate memory for sub-list lengths): %s\n",cudaGetErrorString(error));
    cleanUp(skel);
    return 1;
  }
  
  return 0;
}

extern "C" void cleanUp(Skeletoniser *skel) {
	if (skel->voxels_dev!=NULL) {
		cudaFree(skel->voxels_dev);
		skel->voxels_dev=NULL;
	}
	if (skel->gather_dev!=NULL) {
		for (int i=0;i<=skel->gatherLevels;i++) {
			if (skel->gather_dev[i]!=NULL) {
				cudaFree(skel->gather_dev[i]);
				skel->gather_dev[i]=NULL;
			}
		}
		free(skel->gather_dev);
		skel->gather_dev=NULL;
	}
}

extern "C" void printInfo() {
	size_t freeMem,totalMem, usedMem;
	cuMemGetInfo(&freeMem,&totalMem);
	usedMem=totalMem-freeMem;
	printf("Total memory: %0.2f MB, free: %0.2f MB, used: %0.2f MB\n",
      ((double)totalMem)/1024.0/1024.0,
      ((double)freeMem)/1024.0/1024.0, 
      ((double)usedMem)/1024.0/1024.0);
}

int fetchVolume(Skeletoniser *skel,unsigned char **stepBuffer,int stepCounter) {
  cudaError_t error;
  if (stepBuffer[stepCounter]==0) {
    stepBuffer[stepCounter]=(unsigned char *)malloc(skel->voxelMemSize);
    cudaMemcpy(stepBuffer[stepCounter],skel->voxels_dev,skel->voxelMemSize,cudaMemcpyDeviceToHost);
    error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not copy voxel data from device): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
  }
  return 0;
}

extern "C" int skeletonise(Skeletoniser *skel,unsigned char *voxels,unsigned char **stepBuffer) {
  cudaError_t error;
  /* copy voxels to device memory */
  cudaGetLastError();
  cudaMemcpy(skel->voxels_dev,voxels,skel->voxelMemSize,cudaMemcpyHostToDevice);
  error=cudaGetLastError();
  if (error!=cudaSuccess) {
    printf("CUDA error (could not copy voxel data to device): %s\n",cudaGetErrorString(error));
    cleanUp(skel);
    return 1;
  }

  /* execute the kernels */
  int counter=0;
  unsigned char notFinished=1;

  cudaGetLastError();
  projectVolume_kernel<<<skel->oneDimGrid,skel->oneDimBlock>>>(skel->voxels_dev,skel->volDim,skel->yzPlane_dev);
  cudaThreadSynchronize();
  error=cudaGetLastError();
  if (error!=cudaSuccess) {
    printf("CUDA error (could not execute projectVolume_kernel successfully): %s\n",cudaGetErrorString(error));
    cleanUp(skel);
    return 1;
  }
  findStripCoordinates_kernel<<<skel->findCoordinatesGrid,skel->oneDimBlock>>>(skel->yzPlane_dev,skel->planeDim,skel->coordinateListWork_dev,skel->subListLengths_dev);
  cudaThreadSynchronize();
  error=cudaGetLastError();
  if (error!=cudaSuccess) {
    printf("CUDA error (could not execute findStripCoordinates_kernel successfully): %s\n",cudaGetErrorString(error));
    cleanUp(skel);
    return 1;
  }
  cudaMemcpy(skel->subListLengths,skel->subListLengths_dev,skel->subListLengthsMemSize,cudaMemcpyDeviceToHost);
  
  int numCoordinates=0;
  int numNonZeroSubListLengths=0;
  for (unsigned short i=0;i<skel->planeDim.x*skel->planeDim.y/skel->numThreads;i++) {
    if (skel->subListLengths[i]!=0) {
      numCoordinates+=skel->subListLengths[i];
      skel->nonZeroSubListLengths[numNonZeroSubListLengths]=i;
      numNonZeroSubListLengths++;
    }
  }
  error=cudaMalloc((void **) &(skel->coordinateList_dev),2*numCoordinates*sizeof(unsigned short));
  if (error!=cudaSuccess) {
    printf("CUDA error (could not dynamically allocate memory for strip coordinate list): %s\n",cudaGetErrorString(error));
    cleanUp(skel);
    return 1;
  }
  int destinationIndex=0;
  for (int i=0;i<numNonZeroSubListLengths;i++) {
    int sourceIndex=skel->nonZeroSubListLengths[i]*3*skel->numThreads+1;
    int numShorts=skel->subListLengths[skel->nonZeroSubListLengths[i]]*2;
    cudaMemcpy(&(skel->coordinateList_dev[destinationIndex]),&(skel->coordinateListWork_dev[sourceIndex]),numShorts*sizeof(unsigned short),cudaMemcpyDeviceToDevice);
    destinationIndex+=numShorts;
  }

  int grid = numCoordinates;
  if(grid%2!=0)grid++;
  dim3 block;
  block.x = skel->volDim.x;
  block.y = 2;
  block.z = 1;

  int stepCounter=0;
  while (notFinished) {
    counter++;
    cudaGetLastError();

    findSimpleVoxels_kernel<<<grid/2,block>>>(numCoordinates,skel->voxels_dev,skel->volDim,skel->coordinateList_dev);
    cudaThreadSynchronize();
        error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not execute findCriticalVoxels_kernel successfully): %s\n",cudaGetErrorString(error));
      printInfo();
      cleanUp(skel);
      return 1;
    }

    if (stepBuffer!=0) {
      if (fetchVolume(skel,stepBuffer,stepCounter)==0) stepCounter++;
    }

    cudaGetLastError();
    markSurfaceVoxels_kernel<<<grid/2,block>>>(numCoordinates,skel->voxels_dev,skel->volDim,skel->coordinateList_dev);   
    
    cudaThreadSynchronize();
    error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not execute markSurfaceVoxels_kernel successfully): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
    if (stepBuffer!=0) {
      if (fetchVolume(skel,stepBuffer,stepCounter)==0) stepCounter++;
    }

    cudaGetLastError();
    
   findSurfaceVoxels_kernel<<<grid/2,block>>>(numCoordinates,skel->voxels_dev,skel->volDim,skel->coordinateList_dev);   
   cudaThreadSynchronize();
    error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not execute findSurfaceVoxels_kernel successfully): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
    if (stepBuffer!=0) {
      if (fetchVolume(skel,stepBuffer,stepCounter)==0) stepCounter++;
    }

    cudaGetLastError();
    
    unmarkNonSimpleVoxels_kernel<<<grid/2,block>>>(numCoordinates,skel->voxels_dev,skel->volDim,skel->coordinateList_dev);   
    cudaThreadSynchronize();
    error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not execute unmarkNonSimpleVoxels_kernel successfully): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
    if (stepBuffer!=0) {
      if (fetchVolume(skel,stepBuffer,stepCounter)==0) stepCounter++;
    }

    cudaGetLastError();
    find2CrucialVoxels_kernel<<<grid/2,block>>>(numCoordinates,skel->voxels_dev,skel->volDim,skel->coordinateList_dev);   
    cudaThreadSynchronize();
    error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not execute find2CrucialVoxels_kernel successfully): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
    if (stepBuffer!=0) {
      if (fetchVolume(skel,stepBuffer,stepCounter)==0) stepCounter++;
    }

    cudaGetLastError();
    find1CrucialVoxelsZ_kernel<<<grid/2,block>>>(numCoordinates,skel->voxels_dev,skel->volDim,skel->coordinateList_dev);   
    cudaThreadSynchronize();
    error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not execute find1CrucialVoxels_kernel successfully): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
    cudaGetLastError();
    find1CrucialVoxelsX_kernel<<<grid/2,block>>>(numCoordinates,skel->voxels_dev,skel->volDim,skel->coordinateList_dev);   
    cudaThreadSynchronize();
    error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not execute find1CrucialVoxels_kernel successfully): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
    cudaGetLastError();
    find1CrucialVoxelsY_kernel<<<grid/2,block>>>(numCoordinates,skel->voxels_dev,skel->volDim,skel->coordinateList_dev);   
    cudaThreadSynchronize();
    error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not execute find1CrucialVoxels_kernel successfully): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
    if (stepBuffer!=0) {
      if (fetchVolume(skel,stepBuffer,stepCounter)==0) stepCounter++;
    }

    cudaGetLastError();
    find0CrucialVoxels_kernel<<<grid/2,block>>>(numCoordinates,skel->voxels_dev,skel->volDim,skel->coordinateList_dev);
    cudaThreadSynchronize();
    error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not execute find0CrucialVoxels_kernel successfully): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
    if (stepBuffer!=0) {
      if (fetchVolume(skel,stepBuffer,stepCounter)==0) stepCounter++;
    }

    cudaGetLastError();
    findUnionOfSets<<<skel->oneDimGrid,skel->oneDimBlock>>>(skel->voxels_dev,skel->gather_dev[1],skel->volDim);
    cudaThreadSynchronize();
    error=cudaGetLastError();
    if (error!=cudaSuccess) {
      printf("CUDA error (could not execute findUnionOfSets_kernel successfully): %s\n",cudaGetErrorString(error));
      cleanUp(skel);
      return 1;
    }
    if (stepBuffer!=0) {
      if (fetchVolume(skel,stepBuffer,stepCounter)==0) stepCounter++;
    }
    
    notFinished=0;
    cudaMemcpy(skel->gathered,skel->gather_dev[1],skel->gatherMemSizes[1],cudaMemcpyDeviceToHost);
    for (int i=0;i<skel->oneDimGrid.x;i++) {
      notFinished|=skel->gathered[i];
    }
  }

  /* copy result from device */
  cudaMemcpy(voxels,skel->voxels_dev,skel->voxelMemSize,cudaMemcpyDeviceToHost);
  error=cudaGetLastError();
  if (error!=cudaSuccess) {
    printf("CUDA error (could not copy voxel data from device): %s\n",cudaGetErrorString(error));
    cleanUp(skel);
    return 1;
  }
  return 0;
}
